#include <rknn_api.h>
#include <stdio.h>
#include <errno.h>
#include <cstring>
#include <cstdlib>
#include <unistd.h>
#include <pthread.h>
#include <signal.h>
#include <sys/poll.h>

#include "rk_defines.h"
#include "rk_debug.h"
#include "rk_mpi_ai.h"
#include "rk_mpi_ao.h"
#include "rk_mpi_aenc.h"
#include "rk_mpi_adec.h"
#include "rk_mpi_sys.h"
#include "rk_mpi_mb.h"

// #include </usr/include/eigen3/Eigen/Dense>
#include "test_comm_argparse.h"
#include <vector>
#include "taiic_key.h"
#include <fcntl.h>
#include "stft.h"
#include "taiic_rknn.h"
#include<cmath>

using namespace std;
// using namespace Eigen;

// #ifdef __cplusplus
// #if __cplusplus
// extern "C" {
// #endif
// #endif /* End of #ifdef __cplusplus */

/*RK_U64 TEST_COMM_GetNowUs()
{
    struct timespec time = {0, 0};
    clock_gettime(CLOCK_MONOTONIC, &time);
    return (RK_U64)time.tv_sec * 1000000 + (RK_U64)time.tv_nsec / 1000;  //microseconds
}*/

//四维数组初始化
void arr_initial(RK_U8 ****dt, int dim1, int dim2, int dim3, int dim4)
{
    for (int i = 0; i < dim1; i++)
    {
        for (int j = 0; j < dim2; j++)
        {
            for (int p = 0; p < dim3; p++)
            {
                for (int q = 0; q < dim4; q++)
                {
                    dt[i][j][p][q] = 0;
                }
            }
        }
    }
}

void arr_initialf(float ****dt, int dim1, int dim2, int dim3, int dim4)
{
    for (int i = 0; i < dim1; i++)
    {
        for (int j = 0; j < dim2; j++)
        {
            for (int p = 0; p < dim3; p++)
            {
                for (int q = 0; q < dim4; q++)
                {
                    dt[i][j][p][q] = 0.0;
                }
            }
        }
    }
}

//四维数组内存释放
// template<class T>
void arr_memfree(RK_U8 ****dt, int dim1, int dim2, int dim3, int dim4)
{
    for (int i = 0; i < dim1; i++)
    {
        for (int j = 0; j < dim2; j++)
        {
            for (int p = 0; p < dim3; p++)
            {
                for (int q = 0; q < dim4; q++)
                {
                    free(dt[i][j][p]);
                }
                free(dt[i][j]);
            }
            free(dt[i]);
        }
        free(dt);
    }
}

//四维数组内存释放
void arr_memfreef(float ****dt, int dim1, int dim2, int dim3, int dim4)
{
    for (int i = 0; i < dim1; i++)
    {
        for (int j = 0; j < dim2; j++)
        {
            for (int p = 0; p < dim3; p++)
            {
                for (int q = 0; q < dim4; q++)
                {
                    free(dt[i][j][p]);
                }
                free(dt[i][j]);
            }
            free(dt[i]);
        }
        free(dt);
    }
}

// Hanning window
void hanning(int win_size, float *buffer)
{
    for (int i = 0; i < win_size; i++)
    {
        buffer[i] = (float)0.5 * (1 - cos(i * 2 * PI / (win_size - 1)));
    }
}

// STFT
/*void stft_dccrn(vector<float>& input, RK_U8 ****tmp_out)
{
    //initial
   arr_initial(tmp_out,Model_N,Model_H,Model_W,Model_C);

   RK_LOGI("STFT Starting Test!!! \n");
   assert(chunksize%stftstride==0);
   assert(ovsize%stftstride==0);

   int fftsize = FILTER_LENGTH;//256
   int window_size = WINDOW_LENGTH;//200
   int stride = HOP_LENGTH;//50

   float ****output_arr;
   arr_initialf(output_arr,Model_N,Model_H,Model_WW,Model_C);

   float *window_func=new float[window_size];//window func
   hanning(window_size, window_func);

   for(int i=0;i<chunkframe;i++)
   {
      VectorXf part_input(window_size);
      for(int k=0; k<window_size; k++)
      {
        part_input(k)=input[k+i*stride]* window_func[k];
      }

    //赋值
    MatrixXf fcoef_r(sz,window_size),fcoef_i(sz,window_size);//fcoef_r：实部 fcoef_i：虚部
    for(int w=0; w<sz; w++)
    {
       for(int t=0; t<window_size; t++)
       {
          fcoef_r(w,t) = cos(2. * PI * w * t / fftsize);
          fcoef_i(w,t)= -sin(2. * PI * w * t / fftsize);
       }
    }
    //(129*200) * (200*1)
    MatrixXf fcoef_rr(sz,1),fcoef_ii(sz,1);
    fcoef_rr = fcoef_r * part_input;
    fcoef_ii = fcoef_i * part_input;
    MatrixXf fcoef(sz,2);
    fcoef << fcoef_rr,fcoef_ii;//129*2
    for(int p=0;p<sz;p++)
    {
        for(int q=0;q<fcoef.cols();q++)
        {
            output_arr[0][i][p][q]=fcoef(p,q);
        }
    }
   }
   //129*2->128*2：去掉第一维度，实部/虚部为0
   //transform output
   for(int i=0;i<Model_N;i++)
   {
        for(int j=0;j<Model_H;j++)
        {
            for(int p=0;p>Model_W;p++)
            {
                for(int q=0;q<Model_C;q++)
                {
                    tmp_out[i][j][p][q]=(RK_U8)output_arr[i][j][p+1][q]; //1 157 128 2
                }
            }
        }
    }
    //free memory
    arr_memfreef(output_arr,Model_N,Model_H,Model_WW,Model_C);
    RK_LOGI("Finish STFT\n");
}*/

// ISTFT
/*int istft_dccrn(RK_U8 ****istft_in,RK_U8* tmp_out)
{
    memset(tmp_out,0,Samples);

    int fftsize = FILTER_LENGTH;
    int window_size = WINDOW_LENGTH;
    int stride = HOP_LENGTH;
    int n_samples=(chunkframe-1)*stride+window_size;

    float gain_ifft=(2.0*stride)/window_size;
    float *window_func=new float[window_size];//window func
    hanning(window_size, window_func);
    for(int i=0;i<window_size;i++)
    {
        window_func[i]=gain_ifft*window_func[i];
    }

    MatrixXf coef_cos(sz,window_size),coef_sin(sz,window_size);
    for(int w=0;w<sz;w++)
    {
        double alpha;
        if(w==0 || w==fftsize/2)
        {
            alpha=1.0;
        }
        else
        {
            alpha=2.0;
        }
        alpha /= fftsize;
        for(int t=0;t<window_size;t++)
        {
            coef_cos(w,t) = alpha * cos(2. * PI * w * t / fftsize)*window_func[t];
            coef_sin(w,t) = alpha * sin(2. * PI * w * t / fftsize)*window_func[t];
        }
    }

    MatrixXf tmp_coef_cos(chunkframe,window_size),tmp_coef_sin(chunkframe,window_size);
    MatrixXf tmp_input_cos(chunkframe,sz),tmp_input_sin(chunkframe,sz);
    MatrixXf output_wav(chunkframe,window_size);
    for(int i=0;i<chunkframe;i++)
    {
        for(int j=0;j<sz;j++)
        {
            tmp_input_cos(i,j)=istft_in[0][i][j][0];
            tmp_input_sin(i,j)=istft_in[0][i][j][1];
        }
    }
    tmp_coef_cos=tmp_input_cos*coef_cos;
    tmp_coef_sin=tmp_input_sin*coef_sin;
    output_wav=tmp_coef_cos-tmp_coef_sin;

    for(int i=0;i<chunkframe;i++)
    {
        for(int j=0;j<window_size;j++)
        {
            tmp_out[i*stride+j]=(RK_U8)output_wav(i,j);
        }
    }
    return 0;
}*/

/*--------------------------------------------------------------
                    model inference define
----------------------------------------------------------------*/
// static void dump_tensor_attr(rknn_tensor_attr *attr)
// {
//     char dims[128] = {0};
//     for (int i = 0; i < attr->n_dims; ++i)
//     {
//         int idx = strlen(dims);
//         sprintf(&dims[idx], "%d%s", attr->dims[i], (i == attr->n_dims - 1) ? "" : ", ");
//     }
//     RK_LOGI("  index=%d, name=%s, n_dims=%d, dims=[%s], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
//              "zp=%d, scale=%f\n",
//              attr->index, attr->name, attr->n_dims, dims, attr->n_elems, attr->size, get_format_string(attr->fmt),
//              get_type_string(attr->type), get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
// }

// int NC1HWC2_int8_to_NCHW_int8(const int8_t *src, int8_t *dst, int *dims, int channel, int h, int w)
// {
//     int batch = dims[0];
//     int C1 = dims[1];
//     int C2 = dims[4];
//     int hw_src = dims[2] * dims[3];
//     int hw_dst = h * w;
//     for (int i = 0; i < batch; i++)
//     {
//         src = src + i * C1 * hw_src * C2;
//         dst = dst + i * channel * hw_dst;
//         for (int c = 0; c < channel; ++c)
//         {
//             int plane = c / C2;
//             const int8_t *src_c = plane * hw_src * C2 + src;
//             int offset = c % C2;
//             for (int cur_h = 0; cur_h < h; ++cur_h)
//                 for (int cur_w = 0; cur_w < w; ++cur_w)
//                 {
//                     int cur_hw = cur_h * w + cur_w;
//                     dst[c * hw_dst + cur_h * w + cur_w] = src_c[C2 * cur_hw + offset];
//                 }
//         }
//     }
//     return 0;
// }

static int rknn_toolkit_config(TEST_TOOLKIT_CFG *tkcfg)
{
    RK_LOGI("rknn_toolkit_config\n");
    tkcfg->model_path = "/oem/usr/model/dccrn.rknn";
    // rk_ctx->model_path = "/data/_ckpt_epoch_40_quant.rknn";

    tkcfg->input_type = RKNN_TENSOR_UINT8;
    tkcfg->input_layout = RKNN_TENSOR_NHWC;

    rknn_init(&tkcfg->ctx, tkcfg->model_path, 0, 0, NULL);
    // Get sdk and driver version
    rknn_sdk_version sdk_ver;
    tkcfg->ret = rknn_query(tkcfg->ctx, RKNN_QUERY_SDK_VERSION, &sdk_ver, sizeof(sdk_ver));
    if (tkcfg->ret < 0)
    {
        RK_LOGE("DCCRN:rknn_query fail! ret_ai=%d\n", tkcfg->ret);
        return -1;
    }
    RK_LOGI("rknn_api/rknnrt version: %s, driver version: %s\n", sdk_ver.api_version, sdk_ver.drv_version);
    // get io info
    tkcfg->ret = rknn_query(tkcfg->ctx, RKNN_QUERY_IN_OUT_NUM, &tkcfg->io_num, sizeof(tkcfg->io_num));
    if (tkcfg->ret < 0)
    {
        RK_LOGE("rknn_query fail! ret=%d\n", tkcfg->ret);
        return RK_FAILURE;
    }
    RK_LOGI("------model input num: %d, output num: %d------\n", tkcfg->io_num.n_input, tkcfg->io_num.n_output);

    // input tensor attr
    for (uint32_t i = 0; i < tkcfg->io_num.n_input; i++)
    {
        tkcfg->input_attrs[i].index = i;
        // query info
        tkcfg->ret = rknn_query(tkcfg->ctx, RKNN_QUERY_INPUT_ATTR, &(tkcfg->input_attrs[i]), sizeof(rknn_tensor_attr));
        if (tkcfg->ret < 0)
        {
            RK_LOGE("rknn_init error! ret_ai=%d\n", tkcfg->ret);
            return RK_FAILURE;
        }
        dump_tensor_attr(&tkcfg->input_attrs[i]);
    }

    // output tensor attr
    for (uint32_t i = 0; i < tkcfg->io_num.n_output; i++)
    {
        tkcfg->output_attrs[i].index = i;
        // query info
        tkcfg->ret = rknn_query(tkcfg->ctx, RKNN_QUERY_NATIVE_OUTPUT_ATTR, &(tkcfg->output_attrs[i]), sizeof(rknn_tensor_attr));
        if (tkcfg->ret < 0)
        {
            RK_LOGE("rknn_query fail! ret_ai=%d\n", tkcfg->ret);
            return RK_FAILURE;
        }
        dump_tensor_attr(&tkcfg->output_attrs[i]);
    }

    // Create output tensor memory
    for (uint32_t i = 0; i < tkcfg->io_num.n_output; ++i)
    {
        tkcfg->output_mems[i] = rknn_create_mem(tkcfg->ctx, tkcfg->output_attrs[i].size_with_stride);
    }
    // Set output tensor memory
    for (uint32_t i = 0; i < tkcfg->io_num.n_output; ++i)
    {
        // set output memory and attribute
        tkcfg->ret = rknn_set_io_mem(tkcfg->ctx, tkcfg->output_mems[i], &tkcfg->output_attrs[i]);
        if (tkcfg->ret < 0)
        {
            RK_LOGE("rknn_set_io_mem fail! ret=%d\n", -1);
            return RK_FAILURE;
        }
    }
    // Create input tensor memory
    tkcfg->input_attrs[0].type = tkcfg->input_type;
    tkcfg->input_attrs[0].fmt = tkcfg->input_layout;
    tkcfg->input_mems[0] = rknn_create_mem(tkcfg->ctx, tkcfg->input_attrs[0].size_with_stride); // size
    return RK_SUCCESS;
}

static int rknn_toolkit_input(TEST_TOOLKIT_CFG *tkcfg, RK_U8 *model_input)
{
    RK_LOGI("rknn_toolkit_input\n");
    // Copy input data to input tensor memory
    int width = tkcfg->input_attrs[0].dims[2];
    int stride = tkcfg->input_attrs[0].w_stride;
    if (width == stride)
    {
        memcpy(tkcfg->input_mems[0]->virt_addr, model_input, width * tkcfg->input_attrs[0].dims[1] * tkcfg->input_attrs[0].dims[3]);
    }
    else
    {
        int height = tkcfg->input_attrs[0].dims[1];
        int channel = tkcfg->input_attrs[0].dims[3];
        // copy from src to dst with stride
        RK_U8 *src_ptr = model_input;

        RK_U8 *dst_ptr = (RK_U8 *)tkcfg->input_mems[0]->virt_addr;
        // width-channel elements
        int src_wc_elems = width * channel;
        int dst_wc_elems = stride * channel;
        for (int h = 0; h < height; ++h)
        {
            memcpy(dst_ptr, src_ptr, src_wc_elems);
            src_ptr += src_wc_elems;
            dst_ptr += dst_wc_elems;
        }
    }
    // Set input tensor memory
    tkcfg->ret = rknn_set_io_mem(tkcfg->ctx, tkcfg->input_mems[0], &tkcfg->input_attrs[0]);
    if (tkcfg->ret < 0)
    {
        RK_LOGE("rknn_set_io_mem fail! ret_ai=%d\n", tkcfg->ret);
        return RK_FAILURE;
    }
    return RK_SUCCESS;
}

static int rknn_toolkit_run(TEST_TOOLKIT_CFG *tkcfg)
{
    RK_LOGI("rknn_toolkit_run\n");
    RK_U64 start_us = TEST_COMM_GetNowUs();
    int s32Ret = rknn_run(tkcfg->ctx, NULL);
    if (s32Ret < 0)
    {
        RK_LOGE("dccrn rknn run error: s32Ret= %d\n", s32Ret);
        return -1;
    }
    RK_LOGI("model run:s32Ret = %d", s32Ret);
    RK_U64 elapse_us = TEST_COMM_GetNowUs() - start_us;
    RK_LOGI("------cost time %.2fms------\n", elapse_us / 1000.0f);
}

static int rknn_toolkit_output(TEST_TOOLKIT_CFG *tkcfg)
{
    RK_LOGI("rknn_toolkit_output\n");
    RK_LOGI("=====output origin tensors=========:\n");
    memset(&tkcfg->orig_output_attrs, 0, sizeof(rknn_tensor_attr));

    for (uint32_t i = 0; i < tkcfg->io_num.n_output; i++)
    {
        tkcfg->orig_output_attrs[i].index = i;
        // query info
        tkcfg->ret = rknn_query(tkcfg->ctx, RKNN_QUERY_OUTPUT_ATTR, &(tkcfg->orig_output_attrs[i]), sizeof(rknn_tensor_attr));
        if (tkcfg->ret < 0)
        {
            RK_LOGE("rknn_query fail! ret=%d\n", tkcfg->ret);
            return -1;
        }
        dump_tensor_attr(&tkcfg->orig_output_attrs[i]);
    }

    for (uint32_t i = 0; i < tkcfg->io_num.n_output; ++i)
    {
        int size = tkcfg->orig_output_attrs[i].size_with_stride;
        tkcfg->output_mems_nchw[i] = (int8_t *)malloc(size);
    }

    for (uint32_t i = 0; i < tkcfg->io_num.n_output; i++)
    {
        int channel = tkcfg->orig_output_attrs[i].dims[1];
        int h = tkcfg->orig_output_attrs[i].n_dims > 2 ? tkcfg->orig_output_attrs[i].dims[2] : 1;
        int w = tkcfg->orig_output_attrs[i].n_dims > 3 ? tkcfg->orig_output_attrs[i].dims[3] : 1;
        int hw = h * w;
        RK_LOGI("n_dims=%d,channel=%d,h=%d,w=%d\n", tkcfg->orig_output_attrs[0].n_dims, channel, h, w);
        NC1HWC2_int8_to_NCHW_int8((int8_t *)tkcfg->output_mems[i]->virt_addr, (int8_t *)tkcfg->output_mems_nchw[i], (int *)tkcfg->output_attrs[i].dims,
                                  channel, h, w);
    }

    int model_width = 0;
    int model_height = 0;
    if (tkcfg->input_attrs[0].fmt == RKNN_TENSOR_NCHW)
    {
        RK_LOGI("model is NCHW input fmt\n");
        model_width = tkcfg->input_attrs[0].dims[2];
        model_height = tkcfg->input_attrs[0].dims[3];
    }
    else
    {
        RK_LOGI("model is NHWC input fmt\n");
        model_width = tkcfg->input_attrs[0].dims[1];
        model_height = tkcfg->input_attrs[0].dims[2];
    }
    return RK_SUCCESS;
}

// model inference
static void rknn_model_inference(TEST_TOOLKIT_CFG *tkcfg, RK_U8 *model_input, RK_U64 stmp)
{
    RK_LOGI("----------Model Inference Starting----------\n");
    if (stmp == 1)
    {
        rknn_toolkit_config(tkcfg);
    }
    else
    {
        rknn_toolkit_input(tkcfg, model_input);
        rknn_toolkit_run(tkcfg);
        rknn_toolkit_output(tkcfg);
    }
    RK_LOGI("----------Model Inference Ending----------\n");
}

/*--------------------------------------------------------
                              AIO
----------------------------------------------------------*/
RK_S32 test_ai_poll_event(RK_S32 timeoutMsec, RK_S32 fd)
{
    RK_S32 num_fds = 1;
    struct pollfd pollFds[num_fds];
    RK_S32 ret = 0;

    RK_ASSERT(fd > 0);
    memset(pollFds, 0, sizeof(pollFds));
    pollFds[0].fd = fd;
    pollFds[0].events = (POLLPRI | POLLIN | POLLERR | POLLNVAL | POLLHUP);

    ret = poll(pollFds, num_fds, timeoutMsec);

    if (ret > 0 && (pollFds[0].revents & (POLLERR | POLLNVAL | POLLHUP)))
    {
        RK_LOGE("fd:%d polled error", fd);
        return -1;
    }
    return ret;
}

// #ifdef __cplusplus
// #if __cplusplus
// }
// #endif

// #endif /* End of #ifdef __cplusplus */
